# 导包
import copy
from lxml.etree import Element, SubElement, tostring, ElementTree
 
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
 
 # 类别列表（根据自己的开发需求的实际情况填写）
classes = ['brand', 'cucumber', 'eggplant', 'tomatoes']  #类别

 # label中锚框坐标归一化
def convert(size, box):  # size:(原图w,原图h) , box:(xmin,xmax,ymin,ymax)
    dw = 1./size[0]  # 1./w
    dh = 1./size[1]  # 1./h
    x = (box[0] + box[1])/2.0  # 物体在图中的中心点x坐标
    y = (box[2] + box[3])/2.0  # 物体在图中的中心点y坐标
    w = box[1] - box[0]  # 物体实际像素宽度
    h = box[3] - box[2]  # 物体实际像素高度
    x = x*dw  # 物体中心点x的坐标比（相当于 x/原图w）
    w = w*dw  # 物体宽度的宽度比（相当于 w/原图w）
    y = y*dh # 物体中心点y的坐标比（相当于 y/原图h）
    h = h*dh # 物体高度的高度比（相当于 h/原图h）
    return (x,y,w,h)  # 返回相对于原图的物体中心的（x坐标比，y坐标比，宽度比， 高度比），取值范围[0-1]

# Label格式转化 
def convert_annotation(image_id):

    in_file = open('all_xml/%s.xml'%(image_id)) # 需要转化的标签路径
    out_file = open('to_txt/%s.txt'%(image_id),'w') # 生成txt格式的标签文件（label）的保存路径
    tree=ET.parse(in_file) # 解析xml文件
    root = tree.getroot() # 获取xml文件的根节点
    size = root.find('size')   # 获取指定节点的图像尺寸
    w = int(size.find('width').text)  # 获取图像的宽
    h = int(size.find('height').text) # 获取图像的高
 
    for obj in root.iter('object'):
        cls = obj.find('name').text  # xml里的name参数（类别名称）
        if cls not in classes :
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')   
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
        
        
image_ids_train = open(r'./train.txt').read().strip().split() # 汇总所有.jpg图像文件名称的txt清单（上述生成的）
 
for image_id in image_ids_train:
    print(image_id)
    convert_annotation(image_id) # 转化标注文件格式
